#include <Windows.h>
#include <tchar.h>
#include "resource.h"
#include <Tlhelp32.h>
#pragma comment(lib, "comctl32.lib")

void UnregisterTroyDll(HANDLE hProcess, const TCHAR* szTroyDll)
{
	DWORD dwExitCode;
	GetExitCodeProcess(hProcess, &dwExitCode);
	if (STILL_ACTIVE == dwExitCode)
	{
		const TCHAR* pModuleName = (_tcsrchr(szTroyDll, _T('\\')) + 1);
		MODULEENTRY32 me = {0};
		me.dwSize = sizeof(me);
		HANDLE hSnapShot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, GetProcessId(hProcess));
		if (Module32First(hSnapShot, &me))
		{			
			do 
			{	
				if (0 == _tcsicmp(me.szModule, pModuleName))
				{
					LPVOID fnFreeLibrary = GetProcAddress(GetModuleHandle(_T("Kernel32.dll")), "FreeLibrary");
					CreateRemoteThread(hProcess, NULL, 0, (LPTHREAD_START_ROUTINE)fnFreeLibrary, me.hModule, 0, NULL);
						
					break;
				}
			} while (Module32Next(hSnapShot, &me));
		}
	}
	return;
}

void RegisterTroyDll(HANDLE hProcess, const TCHAR* szTroyDll)
{
	DWORD dwThreadId = 0;
	DWORD dwWritten = 0;
	DWORD dwLen = sizeof(TCHAR) * (_tcslen(szTroyDll) + 1);
	LPVOID pAddr = VirtualAllocEx(hProcess, NULL, dwLen, MEM_COMMIT, PAGE_READWRITE);
	WriteProcessMemory(hProcess, pAddr, szTroyDll, dwLen, &dwWritten);

	LPVOID fnLoadLibrary = GetProcAddress(GetModuleHandle(_T("kernel32.dll")), "LoadLibraryW");
	CreateRemoteThread(hProcess, NULL, 0, (LPTHREAD_START_ROUTINE)fnLoadLibrary, pAddr, 0, &dwThreadId);

	return; 
}

int CALLBACK DlgProc(HWND hDlg, UINT nMsg, WPARAM wParam, LPARAM lParam)
{
	static HANDLE hProcess = NULL;
	static TCHAR szTroyDll[MAX_PATH];
	static PROCESS_INFORMATION pi;
	switch (nMsg)
	{
		case WM_INITDIALOG:
		{
			SetWindowText(hDlg, _T("Demo0124"));

			UINT nWidth		= GetSystemMetrics(SM_CXSCREEN);
			UINT nHeight	= GetSystemMetrics(SM_CYSCREEN);
			RECT rtClient;
			GetClientRect(hDlg, &rtClient);
			SetWindowPos(hDlg, NULL, (nWidth - (rtClient.right - rtClient.left)) / 2, (nHeight - (rtClient.bottom - rtClient.top)) / 2, 0, 0, SWP_NOSIZE);
			
			const TCHAR* szDllName = _T("TroyDll.dll");
			GetModuleFileName(GetModuleHandle(NULL), szTroyDll, MAX_PATH);
			TCHAR* pDir = _tcsrchr(szTroyDll, _T('\\')) + 1;
			*pDir = NULL;
			_tcscat(szTroyDll, szDllName);

			break;
		}
		case WM_COMMAND:
		{
			switch (LOWORD(wParam))
			{
				case IDC_BTNLOAD:
				{
					STARTUPINFO si = {0};
				
					TCHAR szExeName[MAX_PATH] = _T("NotePad.EXE");
					if (CreateProcess(NULL, szExeName, NULL, NULL, FALSE, 0, NULL, NULL, &si, &pi))
					{
						WaitForInputIdle(pi.hProcess, INFINITE);
				
						RegisterTroyDll(pi.hProcess, szTroyDll);
					}
				
					break;
				}
				case IDC_BTNUNLOAD:
				{
					if (NULL != pi.hProcess)
					{
						UnregisterTroyDll(pi.hProcess, szTroyDll);
					}

					break;
				}
			}
			break;
		}
		case WM_CLOSE:
		{
			PostQuitMessage(0);
			break;
		}
	}

	return 0;
}

BOOL WINAPI _tWinMain(HINSTANCE hInstance, HINSTANCE hPreInstance, LPTSTR szCmdLine, int nShowCmd)
{
	DialogBox(hInstance, MAKEINTRESOURCE(IDD_DIALOG), NULL, DlgProc);

	return TRUE;
}